import os
import random
import argparse
from pathlib import Path
import numpy as np


parser = argparse.ArgumentParser(description='.')
parser.add_argument('bin_dir', type=str, help='path to the data preprocessed.')
parser.add_argument('output_dir', type=str, help='path to save output data.')
args = parser.parse_args()

ids_dir = Path(args.bin_dir) / 'input_ids'
mask_dir = Path(args.bin_dir) / 'input_mask'
seg_dir = Path(args.bin_dir) / 'segment_ids'
save_dir = Path(args.output_dir)

batchsize_list = [1, 4, 8, 16, 32, 64]
num_total = len(os.listdir(ids_dir))
for bs in batchsize_list:
	indices = random.sample(list(range(num_total)), bs)
	input_ids, input_mask, segment_ids = [], [], []
	for idx in indices:
		ids_arr = np.fromfile(ids_dir / f"Bert_{idx}.bin", dtype=np.int64)
		input_ids.append(ids_arr.reshape(1, 384))
		mask_arr = np.fromfile(mask_dir / f"input_mask_{idx}.bin", dtype=np.int64)
		input_mask.append(mask_arr.reshape(1, 384))
		seg_arr = np.fromfile(seg_dir / f"segment_ids_{idx}.bin", dtype=np.int64)
		segment_ids.append(seg_arr.reshape(1, 384))
	
	out_ids_dir = save_dir / f'bs{bs}' / 'ids'
	out_ids_dir.mkdir(parents=True, exist_ok=True)
	input_ids = np.vstack(input_ids)
	input_ids.tofile(str(out_ids_dir / 'input_ids.bin'))

	out_mask_dir = save_dir / f'bs{bs}' / 'mask'
	out_mask_dir.mkdir(parents=True, exist_ok=True)
	input_mask = np.vstack(input_mask)
	input_mask.tofile(str(out_mask_dir / 'input_mask.bin'))

	out_seg_dir = save_dir / f'bs{bs}' / 'seg'
	out_seg_dir.mkdir(parents=True, exist_ok=True)
	segment_ids = np.vstack(segment_ids)
	segment_ids.tofile(str(out_seg_dir / 'segment_ids.bin'))
